from flask import Flask, request, Response, jsonify
from flask_cors import CORS  # 导入 Flask-CORS
from stream_predict import predict
import logging
import os
import json

logging.basicConfig(level=logging.INFO,
                    format='%(asctime)s - %(levelname)s - %(message)s')

app = Flask(__name__)
CORS(app)  # 在应用中启用 CORS

DEBUG = os.environ.get('DEBUG', True)  # 从环境变量获取，默认为True
HOST = os.environ.get('HOST', '0.0.0.0')  # 从环境变量获取，默认为0.0.0.0
PORT = int(os.environ.get('PORT', 5000))  # 从环境变量获取，默认为5000

@app.route('/generate', methods=['POST'])
def generate():
    try:
        data = request.get_json()
        if not data or 'prompt' not in data:
            logging.warning("Request body missing 'prompt' field")
            return jsonify({'error': 'Missing prompt field'}), 400

        prompt = data.get('prompt', '')
        if not isinstance(prompt, str):
            logging.warning("Prompt must be a string")
            return jsonify({'error': 'Prompt must be a string'}), 400


        def generate_stream(prompt):
            try:
                # 首先发送包含 prompt 的 JSON 数据块
                yield f"data: {prompt}\n\n"

                for text in predict(prompt):  # 使用封装的 predict 函数
                    yield f"data: {text}\n\n"
            except Exception as e:
                logging.error(f"Error during text generation: {e}")
                yield f"data: {json.dumps({'type': 'error', 'content': str(e)})}\n\n"
            finally:
                logging.info("Streaming completed")

        return Response(generate_stream(prompt), mimetype='text/event-stream')

    except Exception as e:
        logging.exception("Unhandled exception in '/generate' route")
        return jsonify({'error': 'Internal server error'}), 500

if __name__ == '__main__':
    logging.info("Flask app started")
    app.run(debug=DEBUG, host=HOST, port=PORT)
